"""
@Author  : 吕申凯
@Time    : 2022/9/18 14:57
@File    : dropout_regularization.py
@Function: dropout正则化
随机让神经元失活，即随机让某些神经元权重weight=0
减轻了对某个单个的神经元的过度依赖，从而防止过拟合。

注：
    net_normal.eval()为测试环境，在这个环境下，不会近dropout随机失活神经单元。
    而为了让训练环境下得到的数值和测试环境一样多，
    会在训练环境下失活一些神经单元后，将下一层得到的值除上(1-失活概率)，
    这样测试环境和训练环境就会在一个数量级上，方便对比的同时，加快了测试的速度。
"""

import torch
import torch.nn as nn
import matplotlib.pyplot as plt

from torch.utils.tensorboard import SummaryWriter


def gen_data(num_data=10, x_range=(-1, 1)):
    """
        生成数据
    :param num_data:
    :param x_range:
    :return:
    """
    w = 1.5
    train_X = torch.linspace(*x_range).unsqueeze_(1)
    train_Y = w * train_X + torch.normal(0, 0.5, size=train_X.size())
    test_X = torch.linspace(*x_range).unsqueeze_(1)
    test_Y = w * test_X + torch.normal(0, 0.3, size=test_X.size())

    return train_X, train_Y, test_X, test_Y


class MLP(nn.Module):
    """
    模型
    """

    def __init__(self, neural_num, d_prob):
        """
        :param neural_num: 输入维度
        :param d_prob: 正则化随机抛弃神经元的概率
        """
        super(MLP, self).__init__()
        self.linears = nn.Sequential(

            nn.Linear(1, neural_num),
            nn.ReLU(inplace=True),

            nn.Dropout(d_prob),  # 在网络层前面添加正则化
            nn.Linear(neural_num, neural_num),
            nn.ReLU(inplace=True),

            nn.Dropout(d_prob),
            nn.Linear(neural_num, neural_num),
            nn.ReLU(inplace=True),

            nn.Dropout(d_prob),
            nn.Linear(neural_num, 1),
        )

    def forward(self, x):
        """

        :param x: 输入
        :return:
        """
        return self.linears(x)


if __name__ == '__main__':

    torch.manual_seed(1)  # 设置随机种子
    n_hidden = 200  # 输入数据维度
    max_iter = 2000  # 迭代次数
    disp_interval = 400  # 每400次迭代进行一次训练
    lr_init = 0.01  # 学习率

    # 数据
    train_x, train_y, test_x, test_y = gen_data(x_range=(-1, 1))

    # 两个对照实验模型
    net_normal = MLP(neural_num=n_hidden, d_prob=0.)  # 随机抛弃概率为0，则没有使用dropout正则
    net_reglar = MLP(neural_num=n_hidden, d_prob=0.5)  # 随机抛弃概率50%的正则

    # 两个优化器优化器
    optim_normal = torch.optim.SGD(net_normal.parameters(), lr=lr_init, momentum=0.9)
    optim_reglar = torch.optim.SGD(net_reglar.parameters(), lr=lr_init, momentum=0.9)

    # 损失函数
    loss_func = torch.nn.MSELoss()

    # 可视化
    writer = SummaryWriter(comment='_test_tensorboard', filename_suffix="12345678")

    # 迭代训练
    for epoch in range(max_iter):
        # 训练两个模型
        pred_normal, pred_wdecay = net_normal(train_x), net_reglar(train_x)
        loss_normal, loss_wdecay = loss_func(pred_normal, train_y), loss_func(pred_wdecay, train_y)

        optim_normal.zero_grad()
        optim_reglar.zero_grad()

        loss_normal.backward()
        loss_wdecay.backward()

        optim_normal.step()
        optim_reglar.step()

        # 每迭代disp_interval次就进行一次测试
        if (epoch + 1) % disp_interval == 0:

            # 设置为测试环境
            net_normal.eval()
            net_reglar.eval()

            # 可视化
            for name, layer in net_normal.named_parameters():
                writer.add_histogram(name + '_grad_normal', layer.grad, epoch)
                writer.add_histogram(name + '_data_normal', layer, epoch)

            for name, layer in net_reglar.named_parameters():
                writer.add_histogram(name + '_grad_regularization', layer.grad, epoch)
                writer.add_histogram(name + '_data_regularization', layer, epoch)

            test_pred_prob_0, test_pred_prob_05 = net_normal(test_x), net_reglar(test_x)

            # 绘图
            plt.scatter(train_x.data.numpy(), train_y.data.numpy(), c='blue', s=50, alpha=0.3, label='train')
            plt.scatter(test_x.data.numpy(), test_y.data.numpy(), c='red', s=50, alpha=0.3, label='test')
            plt.plot(test_x.data.numpy(), test_pred_prob_0.data.numpy(), 'r-', lw=3, label='d_prob_0')
            plt.plot(test_x.data.numpy(), test_pred_prob_05.data.numpy(), 'b--', lw=3, label='d_prob_05')
            plt.text(-0.25, -1.5, 'd_prob_0 loss={:.8f}'.format(loss_normal.item()),
                     fontdict={'size': 15, 'color': 'red'})
            plt.text(-0.25, -2, 'd_prob_05 loss={:.6f}'.format(loss_wdecay.item()),
                     fontdict={'size': 15, 'color': 'red'})

            plt.ylim((-2.5, 2.5))
            plt.legend(loc='upper left')
            plt.title("Epoch: {}".format(epoch + 1))
            plt.show()
            plt.close()

            # 切换为训练模式
            net_normal.train()
            net_reglar.train()
